#include <fstream>
#include <cmath>
#include "gc.h"
#include "itkImageRegionConstIteratorWithIndex.h"

inline double calcAreaWeight (int A0, int A1, int A01, double sigma, 
			      double& O01)
{
  O01 = A0 > A1 ? double(A01) / double (A0) : double(A01) / double(A1);
  return exp(-(1.0 - O01) * (1.0 - O01) / (2.0 * sigma * sigma));
}



inline double calcCorrelationWeight (double correlation, double sigma)
{
  return exp(-(1.0 - correlation) * (1.0 - correlation) / (2.0 * sigma * sigma));
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
void addEdgeToGraph (Volume const& volume, Graph3D& graph, 
		     int layerNo, int prevLayerNo, 
		     std::string const& correlationPath, 
		     std::string const& areaPath, 
		     double sigma_area, double sigma_correlation, 
		     double alpha = 1.0)
{
  double weight_threshold = 1e-20;
  int layerInc = layerNo - prevLayerNo;
  CorrespondenceMatrix crrMat(correlationPath);
  int componentNum0 = volume.layers[prevLayerNo].components.size();
  int componentNum1 = volume.layers[layerNo].components.size();
  int totalLength = componentNum0 * componentNum1;
  int* areaMat = new int[totalLength];
  readArrayBinary(areaMat, totalLength, areaPath.c_str());
  for (int nodeNo0 = 0; nodeNo0 < graph.nodes[prevLayerNo].size(); nodeNo0++) {
    int cn0 = graph.nodes[prevLayerNo][nodeNo0].componentNo;
    for (int nodeNo1 = 0; nodeNo1 < graph.nodes[layerNo].size(); nodeNo1++) {
      int cn1 = graph.nodes[layerNo][nodeNo1].componentNo;
      int area = areaMat[cn0 * componentNum1 + cn1];
      const Correspondence* crr = 
	crrMat.getCorrespondence(cn0, prevLayerNo, cn1, layerNo);
      if (area > 0 && crr != NULL) {
      // if (area > 0) {
      // if (crr != NULL) {
	double O;
	double weight_area = 
	  calcAreaWeight(volume.layers[prevLayerNo].components[cn0].points.size(), 
			 volume.layers[layerNo].components[cn1].points.size(), 
			 area, sigma_area, O);
	// double correlation = 0.0;
	double correlation = crr->getNorm();
	double weight_correlation = 
	  calcCorrelationWeight(correlation, sigma_correlation);
	double weight = weight_area * weight_correlation * alpha;
	// double weight = (weight_area + weight_correlation) * alpha;
	// double weight = weight_area * alpha;
	// double weight = weight_correlation * alpha;
	if (weight > weight_threshold) {
	  Edge edge(nodeNo1, area, weight_area, 
		    correlation, weight_correlation, weight);
	  if (layerInc == 1) {
	    graph.addEdge(prevLayerNo, nodeNo0, edge);
	  }
	  else {
	    graph.addSkipEdge(prevLayerNo, nodeNo0, edge);
	  }
	}
      }
    }
  }
  if (areaMat != NULL) {
    delete[] areaMat;
  }
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
void generateGraph (Volume const& volume, Graph3D& graph, 
		    std::vector<std::string> const& adjacentCorrelationPaths, 
		    std::vector<std::string> const& skipCorrelationPaths, 
		    std::vector<std::string> const& adjacentAreaPaths, 
		    std::vector<std::string> const& skipAreaPaths, 
		    double sigma_area, double sigma_correlation, double alpha, 
		    VerboseLevel verbose)
{
  int layerNum = volume.layers.size();
  graph.nodes.reserve(layerNum);
  // Add nodes
#ifdef USE_OMP
#pragma omp parallel for
#endif
  for (int layerNo = 0; layerNo < layerNum; layerNo++) {
    int componentNum = volume.layers[layerNo].components.size();
    int nodeNum = componentNum;
    GraphLayerNodeSet graphLayerNodeSet;
    graph.nodes.push_back(graphLayerNodeSet);
    graph.nodes[layerNo].reserve(nodeNum);
    int nodeNo = 0;
    for (int componentNo = 0; componentNo < componentNum; componentNo++) {
      NodeStatus status = volume.layers[layerNo].components[componentNo].isOnBoundary? 
	NODE_DELETED: NODE_NORMAL;
      Node node(nodeNo++, layerNo, componentNo, status);
      graph.addNode(node);
    }
    if (verbose == VERBOSE_HIGH) {
      std::cout << "Generate graph: adding nodes for layer " 
		<< layerNo << " done." << std::endl;
    }
  }
  if (verbose >= VERBOSE_MEDIUM) {
    std::cout << "Generate graph: adding nodes done." << std::endl;
  }
  // Add edges
#ifdef USE_OMP
#pragma omp parallel for
#endif
  for (int layerNo = 1; layerNo < layerNum; layerNo++) {
    // For adjacent slice
    int prevLayerNo = layerNo - 1;
    addEdgeToGraph(volume, graph, layerNo, prevLayerNo, 
		   adjacentCorrelationPaths[prevLayerNo], 
		   adjacentAreaPaths[prevLayerNo], 
		   sigma_area, sigma_correlation);
    if (verbose == VERBOSE_HIGH) {
      std::cout << "Generate graph: adding edges between layer " << prevLayerNo
		<< " & " << layerNo << " done." << std::endl;
    }
    // For skip slice
    if (prevLayerNo > 0) {
      prevLayerNo--;
      addEdgeToGraph(volume, graph, layerNo, prevLayerNo, 
		     skipCorrelationPaths[prevLayerNo], 
		     skipAreaPaths[prevLayerNo], 
		     sigma_area, sigma_correlation, alpha);
      if (verbose == VERBOSE_HIGH) {
	std::cout << "Generate graph: adding skip edges between layer " 
		  << prevLayerNo << " & " << layerNo << " done." << std::endl;
      }
    }
  }
}
